package org.jetlinks.pro.openapi.interceptor;

import org.springframework.core.io.buffer.DataBuffer;
import org.springframework.core.io.buffer.DataBufferUtils;
import org.springframework.http.server.reactive.ServerHttpRequest;
import org.springframework.http.server.reactive.ServerHttpRequestDecorator;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;

import javax.annotation.Nonnull;

public class OpenApiServerHttpRequestDecorator extends ServerHttpRequestDecorator {

   private final Flux<DataBuffer> bufferFlux;

    public OpenApiServerHttpRequestDecorator(ServerHttpRequest delegate) {
        super(delegate);
        bufferFlux = Flux
            .usingWhen(DataBufferUtils.join(delegate.getBody()).cache(),
                       dataBuffer -> Mono.just(DataBufferUtils.retain(dataBuffer)),
                       dataBuffer -> {
                           DataBufferUtils.release(dataBuffer);
                           return Mono.empty();
                       }
            );
    }

    @Nonnull
    @Override
    public Flux<DataBuffer> getBody() {
        return bufferFlux;
    }

}
